import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 确保你的计算机上已经安装了NVIDIA驱动和CUDA工具包，并且PyTorch版本与CUDA版本兼容
# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')


# 定义简单的神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x


# 加载模型
def load_model():
    model = SimpleNN().to(device)  # 创建模型实例并移动到设备上
    model.load_state_dict(torch.load('./build/mnist_simplenn.pth', weights_only=True))  # 加载状态字典
    model.eval()  # 设置模型为评估模式
    return model


# 修改模型参数或权重
def modify_model(model):
    # 例如，修改第一个全连接层的权重
    with torch.no_grad():
        model.fc1.weight += 0.1  # 增加所有权重的值
        model.fc1.bias -= 0.1  # 减少所有偏置的值
    return model


# 评估模型
def evaluate_model(model, test_loader):
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)  # 将数据移动到设备上
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Accuracy of the network on the test images: {100 * correct / total}%')
    return correct / total


# 保存模型
def save_model(model, path='./build/modified_mnist_simplenn.pth'):
    torch.save(model.state_dict(), path)
    print(f'Model saved to {path}')


# 主程序入口
if __name__ == '__main__':
    # 加载已保存的模型
    model = load_model()

    # 假设你有一个名为 `test_loader` 的数据加载器
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False, download=True, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=64, shuffle=True)

    # 修改模型参数或权重
    modified_model = modify_model(model)

    # 评估修改后的模型
    accuracy = evaluate_model(modified_model, test_loader)

    # 保存修改后的模型
    save_model(modified_model)
